# -*- coding: utf-8 -*-
"""
@author: Yonghao.Xu
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import scipy.io as sio  
import numpy as np  
import matplotlib.pyplot as plt
from HyperFunctions import*
import  keras
import time
from keras.datasets import mnist
from keras.models import Model,Sequential,save_model,load_model
from keras.layers import Input, Dense, Activation,LSTM,merge,Conv2D, MaxPooling2D,AveragePooling2D, Flatten,Dropout
from keras.optimizers import SGD
from keras.regularizers import l2
from keras.optimizers import RMSprop
from keras import backend as K
from keras.utils import np_utils
from sklearn.metrics import cohen_kappa_score
import cv2
import scipy.io as scio
import tensorflow as tf

def stats_graph(graph):
    flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
    params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
    print('FLOPs: {};    Trainable params: {}'.format(flops.total_float_ops, params.total_parameters))


def LSTMMCNN_RS(time_step,nb_features,num_PC,img_rows,img_cols):
    LSTMInput = Input(shape=(time_step,nb_features),name='LSTMInput')

    LSTMSpectral = LSTM(128,name='LSTMSpectral',consume_less='gpu',W_regularizer=l2(0.0001),U_regularizer=l2(0.0001))(LSTMInput)
    
    LSTMDense = Dense(128,activation='relu', name='LSTMDense')(LSTMSpectral)
   
    LSTMSOFTMAX = Dense(nb_classes,activation='softmax', name='LSTMSOFTMAX')(LSTMDense)    
    
    CNNInput = Input(shape=[img_rows,img_cols,num_PC],name='CNNInput')

    #print(CNNInput.shape)

    CONV1 = Conv2D(32, (3, 3), strides=1,activation='relu', padding='same', name='CONV1')(CNNInput)
    POOL1 = MaxPooling2D((2, 2), strides=(2,2), name='POOL1')(CONV1)

    #print(POOL1.shape)

    CONV2 = Conv2D(32, (3, 3), strides=1,activation='relu', padding='same', name='CONV2')(POOL1)
    POOL2 = MaxPooling2D((2, 2), strides=(2,2), name='POOL2')(CONV2)

    #print(POOL2.shape)

    CONV3 = Conv2D(32, (3, 3), strides=1,activation='relu', padding='same', name='CONV3')(POOL2)
    #print(CONV3.shape)
    POOL3 = MaxPooling2D((2,2), strides=(2,2), name='POOL3')(CONV3)
    #print(POOL3.shape)

    FLATTEN1 = Flatten(name='FLATTEN1')(POOL1)
    FLATTEN2 = Flatten(name='FLATTEN2')(POOL2)
    FLATTEN3 = Flatten(name='FLATTEN3')(POOL3)
    
    DENSE1 = Dense(128,activation='relu', name='DENSE1')(FLATTEN1)
    DENSE2 = Dense(128,activation='relu', name='DENSE2')(FLATTEN2)
    DENSE3 = Dense(128,activation='relu', name='DENSE3')(FLATTEN3)
    
    CNNDense = keras.layers.Add()([DENSE1, DENSE2, DENSE3])

    
    CNNSOFTMAX = Dense(nb_classes,activation='softmax', name='CNNSOFTMAX')(CNNDense)
    

    JOINT = keras.layers.Concatenate(axis=-1)([LSTMDense,CNNDense])
    JOINTDENSE = Dense(128,activation='relu', name='JOINTDENSE')(JOINT)

    JOINTSOFTMAX = Dense(nb_classes,activation='softmax',name='JOINTSOFTMAX')(JOINTDENSE)


    model = Model(input=[LSTMInput,CNNInput], output=[JOINTSOFTMAX,LSTMSOFTMAX,CNNSOFTMAX])    
    
    
    rmsp = RMSprop(lr=0.001, rho=0.9, epsilon=1e-05)
    
    model.compile(optimizer=rmsp, loss='categorical_crossentropy',
                  metrics=['accuracy'],loss_weights=[1, 1,1])

    return model
    
    
def LSTM_RS(time_step,nb_features):
    LSTMInput = Input(shape=(time_step,nb_features),name='LSTMInput')

    LSTMSpectral = LSTM(128,name='LSTMSpectral',consume_less='gpu',W_regularizer=l2(0.0001),U_regularizer=l2(0.0001))(LSTMInput)
    
    LSTMDense = Dense(128,activation='relu', name='LSTMDense')(LSTMSpectral)
    LSTMSOFTMAX = Dense(nb_classes,activation='softmax', name='LSTMSOFTMAX')(LSTMDense)
    
    model = Model(input=[LSTMInput], output=[LSTMSOFTMAX])
    rmsp = RMSprop(lr=0.001, rho=0.9, epsilon=1e-05)
    
    model.compile(optimizer=rmsp, loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model
    
    
def MCNN_RS(num_PC,img_rows,img_cols):
    CNNInput = Input(shape=[img_rows,img_cols,num_PC],name='CNNInput')

    #print(CNNInput.shape)
    CONV1 = Conv2D(32, (3, 3), activation='relu', border_mode='same', name='CONV1')(CNNInput)
    #print(CONV1.shape)
    POOL1 = MaxPooling2D((2, 2), strides=(2,2), name='POOL1')(CONV1)
    CONV2 = Conv2D(32, (3, 3), activation='relu', border_mode='same', name='CONV2')(POOL1)
    #print(CONV2.shape)
    POOL2 = MaxPooling2D((2, 2), strides=(2,2), name='POOL2')(CONV2)
    CONV3 = Conv2D(32, (3, 3), activation='relu', border_mode='same', name='CONV3')(POOL2)
    #print(CONV3.shape)
    POOL3 = MaxPooling2D((2,2), strides=(2,2), name='POOL3')(CONV3)
    
    FLATTEN1 = Flatten(name='FLATTEN1')(POOL1)
    FLATTEN2 = Flatten(name='FLATTEN2')(POOL2)
    FLATTEN3 = Flatten(name='FLATTEN3')(POOL3)
    
    DENSE1 = Dense(128,activation='relu', name='DENSE1')(FLATTEN1)
    DENSE2 = Dense(128,activation='relu', name='DENSE2')(FLATTEN2)
    DENSE3 = Dense(128,activation='relu', name='DENSE3')(FLATTEN3)
    
    CNNDense = CNNDense = keras.layers.Add()([DENSE1, DENSE2, DENSE3])

    
    CNNSOFTMAX = Dense(nb_classes,activation='softmax', name='CNNSOFTMAX')(CNNDense)
    
    
    model = Model(input=[CNNInput], output=[CNNSOFTMAX])
    rmsp = RMSprop(lr=0.001, rho=0.9, epsilon=1e-05)
    
    model.compile(optimizer=rmsp, loss='categorical_crossentropy',
                  metrics=['accuracy'])

    return model

# Spectral
# dataID=4
#
# if dataID==1:
#     cls_num=9
# elif dataID==2:
#     cls_num=16
# elif dataID==3:
#     cls_num=13
# else:
#     cls_num=15
# w=2
# num_PC=2
# israndom=True
# randtime = 10
# OASpectral_Pavia = np.zeros((cls_num+5,randtime+2))
# s1s2=2
# time_step = 3
#
# data_dic={1:'pavia',2:'indian',3:'ksc',4:'houston'}
#
# for r in range(0,randtime):
#
#     #################Pavia#################
#
#     data = HyperspectralSamples(dataID=dataID, timestep=time_step, w=w, num_PC=num_PC, israndom=israndom, s1s2=s1s2)
#     X = data[0]
#     X_train = data[1]
#     X_test = data[2]
#     XP = data[3]
#     XP_train = data[4]
#     XP_test = data[5]
#     Y = data[6]-1
#     Y_train = data[7]-1
#     Y_test = data[8]-1
#
#     batch_size = 64
#
#     nb_classes = Y_train.max()+1
#     nb_epoch = 500
#     nb_features = X.shape[-1]
#
#     img_rows, img_cols = XP.shape[-1],XP.shape[-1]
#     # convert class vectors to binary class matrices
#     y_train = np_utils.to_categorical(Y_train, nb_classes)
#     y_test = np_utils.to_categorical(Y_test, nb_classes)
#
#     model = LSTM_RS(time_step=time_step,nb_features=nb_features)
#     trn_time1 = time.time()
#     histloss=model.fit([X_train], [y_train], nb_epoch=nb_epoch, batch_size=batch_size, verbose=0, shuffle=True)
#     trn_time2 = time.time()
#     losses = histloss.history
#
#     tes_time1 = time.time()
#     PredictLabel = model.predict([X_test],verbose=0).argmax(axis=-1)
#     tes_time2 = time.time()
#
#     OA,Kappa,ProducerA = CalAccuracy(PredictLabel,Y_test[:,0])
#     AA=np.mean(ProducerA)
#     OASpectral_Pavia[0, r] = OA * 100
#     OASpectral_Pavia[1, r] = AA * 100
#     OASpectral_Pavia[2, r] = Kappa * 100
#     OASpectral_Pavia[3, r] = trn_time2 - trn_time1
#     OASpectral_Pavia[4, r] = tes_time2 - tes_time1
#     OASpectral_Pavia[5:,r] = ProducerA*100
#
#     print('rand',r+1,'LSTM '+ str(data_dic[dataID]) +' test accuracy:', OA*100)
#
#     Map = model.predict([X],verbose=0)
#
#     Spectral = Map.argmax(axis=-1)+1
#
#     #X_result = DrawResult(Spectral,1)
#
#     #plt.imsave('LSTM_Pavia_r'+repr(r+1)+'OA_'+repr(int(OA*10000))+'.png',X_result)
# OASpectral_Pavia[:,-2]=np.mean(OASpectral_Pavia[:,0:-2],axis=1)
# OASpectral_Pavia[:,-1]=np.std(OASpectral_Pavia[:,0:-2],axis=1)
#
# scio.savemat('lstm_strategy_'+str(s1s2)+'_'+str(data_dic[dataID])+'.mat',{'data':OASpectral_Pavia})
#
# if dataID == 1:
#     row = 610
#     col = 340
# elif dataID == 2:
#     row = 145
#     col = 145
# elif dataID == 3:
#     row = 512
#     col = 614
# elif dataID==4:
#     row = 349
#     col = 1905
#
# cv2.imwrite('lstm_strategy_'+str(s1s2)+'_all_' + str(data_dic[dataID]) + '.png', Spectral.reshape(row, col))


#

#Spatial

# dataID=4
#
# if dataID==1:
#     cls_num=9
# elif dataID==2:
#     cls_num=16
# elif dataID==3:
#     cls_num=13
# else:
#     cls_num=15
#
# w=28
# num_PC=4
# israndom=True
# randtime = 10
# OASpatial_Pavia = np.zeros((cls_num+5,randtime+2))
# s1s2=1
# time_step=1
#
# data_dic={1:'pavia',2:'indian',3:'ksc',4:'houston'}
#
# for r in range(0,randtime):
#
#     #################Pavia#################
#     data = HyperspectralSamples(dataID=dataID, timestep=time_step, w=w, num_PC=num_PC, israndom=israndom, s1s2=s1s2)
#     X = data[0]
#     X_train = data[1]
#     X_test = data[2]
#     XP = data[3]
#     XP_train = data[4]
#     XP_test = data[5]
#     Y = data[6]-1
#     Y_train = data[7]-1
#     Y_test = data[8]-1
#
#     batch_size = 64
#
#     nb_classes = Y_train.max()+1
#     nb_epoch = 500
#     nb_features = X.shape[-1]
#
#     img_rows, img_cols = XP.shape[-2],XP.shape[-2]
#     # convert class vectors to binary class matrices
#     y_train = np_utils.to_categorical(Y_train, nb_classes)
#     y_test = np_utils.to_categorical(Y_test, nb_classes)
#
#     model = MCNN_RS(num_PC,img_rows,img_cols)
#
#     trn_time1 = time.time()
#     histloss=model.fit([XP_train], [y_train], nb_epoch=nb_epoch, batch_size=batch_size, verbose=0, shuffle=True)
#     trn_time2 = time.time()
#     losses = histloss.history
#
#     tes_time1 = time.time()
#     PredictLabel = model.predict([XP_test],verbose=0).argmax(axis=-1)
#     tes_time2 = time.time()
#
#     OA,Kappa,ProducerA = CalAccuracy(PredictLabel,Y_test[:,0])
#     AA = np.mean(ProducerA)
#     OASpatial_Pavia[0, r] = OA * 100
#     OASpatial_Pavia[1, r] = AA * 100
#     OASpatial_Pavia[2, r] = Kappa * 100
#     OASpatial_Pavia[3, r] = trn_time2 - trn_time1
#     OASpatial_Pavia[4, r] = tes_time2 - tes_time1
#     OASpatial_Pavia[5:, r] = ProducerA * 100
#
#     print('rand',r+1,'MCNN '+ str(data_dic[dataID]) +' test accuracy:', OA*100)
#
#     Map = model.predict([XP],verbose=0)
#
#     Spatial = Map.argmax(axis=-1)+1
#
#     #X_result = DrawResult(Spatial,1)
#
#     #plt.imsave('MCNN_Pavia_r'+repr(r+1)+'OA_'+repr(int(OA*10000))+'.png',X_result)
#
# OASpatial_Pavia[:,-2]=np.mean(OASpatial_Pavia[:,0:-2],axis=1)
# OASpatial_Pavia[:,-1]=np.std(OASpatial_Pavia[:,0:-2],axis=1)
#
# scio.savemat('mscnn_'+str(data_dic[dataID])+'.mat',{'data':OASpatial_Pavia})
#
# if dataID == 1:
#     row = 610
#     col = 340
# elif dataID == 2:
#     row = 145
#     col = 145
# elif dataID == 3:
#     row = 512
#     col = 614
# elif dataID==4:
#     row = 349
#     col = 1905
#
# cv2.imwrite('mscnn_all_' + str(data_dic[dataID]) + '.png', Spatial.reshape(row, col))


#Joint
dataID=4

if dataID==1:
    cls_num=9
elif dataID==2:
    cls_num=16
elif dataID==3:
    cls_num=13
else:
    cls_num=15

time_step=3
w=28
num_PC=4
israndom=True
s1s2=2
randtime = 10
OAJoint = np.zeros((cls_num+5,randtime+2))

data_dic={1:'pavia',2:'indian',3:'ksc',4:'houston'}

nb_classes=16
#
# for r in range(0,randtime):
#
#     #################Pavia#################
#
#     data = HyperspectralSamples(dataID=dataID, timestep=time_step, w=w, num_PC=num_PC, israndom=israndom, s1s2=s1s2)
#     X = data[0]
#     X_train = data[1]
#     X_test = data[2]
#     XP = data[3]
#     XP_train = data[4]
#     XP_test = data[5]
#     Y = data[6]-1
#     Y_train = data[7]-1
#     Y_test = data[8]-1
#
#     batch_size = 64
#
#     nb_classes = Y_train.max()+1
#     nb_epoch = 500
#     nb_features = X.shape[-1]
#
#     img_rows, img_cols = XP.shape[-2],XP.shape[-2]
#     # convert class vectors to binary class matrices
#     y_train = np_utils.to_categorical(Y_train, nb_classes)
#     y_test = np_utils.to_categorical(Y_test, nb_classes)
#
#
#     model = LSTMMCNN_RS(time_step=time_step,nb_features=nb_features,num_PC=num_PC,img_rows=img_rows,img_cols=img_cols)
#     trn_time1 = time.time()
#     histloss=model.fit([X_train,XP_train], [y_train,y_train,y_train], nb_epoch=nb_epoch, batch_size=batch_size, verbose=0, shuffle=True)
#     trn_time2 = time.time()
#     losses = histloss.history
#
#     tes_time1 = time.time()
#     PredictLabel = model.predict([X_test,XP_test],verbose=0)[0].argmax(axis=-1)
#     tes_time2 = time.time()
#
#     OA,Kappa,ProducerA = CalAccuracy(PredictLabel,Y_test[:,0])
#     AA = np.mean(ProducerA)
#     OAJoint[0, r] = OA * 100
#     OAJoint[1, r] = AA * 100
#     OAJoint[2, r] = Kappa * 100
#     OAJoint[3, r] = trn_time2 - trn_time1
#     OAJoint[4, r] = tes_time2 - tes_time1
#     OAJoint[5:, r] = ProducerA * 100
#
#     print('rand',r+1,'LSTMMCNN_ '+str(data_dic[dataID])+' test accuracy:', OA*100)
#
#     Map = model.predict([X,XP],verbose=0)
#
#     Joint = Map[0].argmax(axis=-1)+1
#
#     #print(Joint.shape)
#
#     #X_result = DrawResult(Joint,1)
#
#     #plt.imsave('LSTMMCNN_Pavia_r'+repr(r+1)+'OA_'+repr(int(OA*10000))+'.png',X_result)
#
# #save_fn = 'OAAll.mat'
# #sio.savemat(save_fn, {'OASpectral_Pavia': OASpectral_Pavia,'OASpatial_Pavia': OASpatial_Pavia,'OAJoint_Pavia': OAJoint_Pavia})
# OAJoint[:,-2]=np.mean(OAJoint[:,0:-2],axis=1)
# OAJoint[:,-1]=np.std(OAJoint[:,0:-2],axis=1)
#
# sio.savemat('SSUN_'+str(data_dic[dataID])+'.mat', {'data': OAJoint})
#
# if dataID == 1:
#     row = 610
#     col = 340
# elif dataID == 2:
#     row = 145
#     col = 145
# elif dataID == 3:
#     row = 512
#     col = 614
# elif dataID==4:
#     row = 349
#     col = 1905
#
# cv2.imwrite('ssun_all_' + str(data_dic[dataID]) + '.png', Joint.reshape(row, col))
#
# print('结果图生成完成！！！！！！！')

if __name__ == '__main__':

    model = LSTMMCNN_RS(time_step=3, nb_features=200, num_PC=4, img_rows=28, img_cols=28)
    sess = K.get_session()
    graph = sess.graph
    stats_graph(graph)

